#pragma once

#include "BaseDetection.h"
#include "NvInfer.h"
#include "PedestrianBox.h"

using namespace nvinfer1;

class YoloV5Detection : public BaseDetection
{
public:
	YoloV5Detection(int device_id, uint8_t dataType, bool use_dynamic=false, bool alway_serialize=false);
	~YoloV5Detection();

	virtual int init(const std::string& weight_path);
	int predict(const std::vector<cv::Mat>& vec_imgs, std::vector<std::vector<Detection>>& final_res, int predict_h = 512, int predict_w = 512);
	int free();

private:
	int doInference(const float* input, float* output, int batchSize, int input_h, int input_w);
	int DeserializeEngine(const std::string& engine_name);
	int init_buffers(int batchSize, int input_h, int input_w);
	int free_buffers();

	IRuntime* runtime;
	ICudaEngine* engine;
	IExecutionContext* context;
	cudaStream_t stream;

	void* device_buffers[2];
	void* host_buffers[2];
	bool alway_serialize;
	bool use_dynamic;
	uint8_t dataType;
	int device_id;
    std::vector<std::string> class_names{"head", "body", "arm_front", "arm_left", "body_back", "arm_right"};
};
